import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from transformers import BertTokenizer
from datasets import load_dataset
import json

import math

# --------------------------
# 配置类 (保持原结构)
# --------------------------
class GPTConfig:
    block_size: int = 1024    # 最大输入序列长度
    batch_size: int = 12     # 训练批量大小
    n_layer: int = 6         # Transformer层数
    n_head: int = 16         # 注意力头数
    n_embd: int = 1024        # 嵌入维度（hidden_dim）
    head_size: int = n_embd // n_head  # 每个注意力头的维度
    dropout: float = 0.1     # Dropout概率
    vocab_size: int = 21128   # 词表大小（GPT-2标准）

# --------------------------
# 注意力机制模块
# --------------------------
class MultiHeadAttention(nn.Module):
    """ 标准的多头注意力实现（更高效版本） """
    def __init__(self, config):
        super().__init__()
        self.n_embd = config.n_embd
        self.n_head = config.n_head
        self.head_size = config.head_size
        
        # 合并的QKV线性投影
        self.qkv = nn.Linear(config.n_embd, 3*config.n_embd)
        self.proj = nn.Linear(config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        
        # 注册因果掩码（无需梯度）
        self.register_buffer("mask", torch.tril(torch.ones(config.block_size, config.block_size)))

    def forward(self, x):
        B, T, C = x.shape  # 批量大小, 序列长度, 特征维度
        
        # 生成QKV并分割成多头
        qkv = self.qkv(x)
        q, k, v = qkv.split(self.n_embd, dim=2)  # 每个形状为(B, T, n_embd)
        
        # 重塑为多头形状 (B, num_heads, T, head_size)
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        
        # 计算注意力权重 (缩放点积注意力)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_size))
        att = att.masked_fill(self.mask[:T, :T] == 0, float('-inf'))  # 因果掩码
        att = F.softmax(att, dim=-1)
        att = self.dropout(att)
        
        # 应用注意力到V值
        y = att @ v  # (B, nh, T, hs)
        
        # 合并多头输出
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # (B, T, C)
        y = self.proj(y)
        return y

# --------------------------
# 前馈神经网络
# --------------------------
class FeedForward(nn.Module):
    """ 两层MLP结构 """
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),  # GPT使用的激活函数
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout)
        )
    
    def forward(self, x):
        return self.net(x)

# --------------------------
# Transformer块
# --------------------------
class Block(nn.Module):
    """ 完整的Transformer块（残差连接 + LayerNorm） """
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = MultiHeadAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ffn = FeedForward(config)

    def forward(self, x):
        # 注意力子层
        x = x + self.attn(self.ln1(x))  # 残差连接
        # 前馈子层
        x = x + self.ffn(self.ln2(x))   # 残差连接
        return x

# --------------------------
# 完整GPT模型
# --------------------------
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config  # 保存配置参数
        
        # 输入嵌入层
        self.token_embed = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_embed = nn.Embedding(config.block_size, config.n_embd)
        
        # Transformer层堆叠
        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
        
        # 输出层
        self.ln_final = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # 权重绑定（输入嵌入与输出层共享权重）
        self.token_embed.weight = self.head.weight
        
        # 初始化权重
        self.apply(self._init_weights)
        
    def _init_weights(self, module):
        """ 初始化策略 """
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        
        # 生成位置索引（自动处理序列截断）
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)  # (T)
        
        # 嵌入求和（词嵌入 + 位置嵌入）
        tok_emb = self.token_embed(idx)  # (B,T,C)
        pos_emb = self.pos_embed(pos)    # (T,C)
        x = tok_emb + pos_emb  # (B,T,C)
        
        # 通过Transformer块
        x = self.blocks(x)
        x = self.ln_final(x)
        logits = self.head(x)  # (B,T,vocab_size)
        
        # 计算损失
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), 
                targets.view(-1),
                ignore_index=-1  # 可设置忽略特定索引
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        """ 自回归文本生成 """
        for _ in range(max_new_tokens):
            # 截断输入到最大允许长度
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            
            # 前向计算
            logits, _ = self(idx_cond)
            
            # 聚焦最后一个时间步
            logits = logits[:, -1, :]  # (B, vocab_size)
            
            # 概率采样
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
            
            # 拼接新token
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)
        return idx

def train(model,optimizer,scheduler,train_loader,val_loader,device):
    model.train()
    total_loss=0
    for batch_idx,(x,y) in enumerate(train_loader):
        x,y = x.to(device),y.to(device)
        logits,loss = model(x,targets=y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()
        total_loss +=loss.item()
        if batch_idx % 100 ==0:
            print(f'Batch: {batch_idx}, Loss: {loss.item():.4f}')

    return total_loss

def eval(model, val_loader, device):
    # 验证
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            logits, loss = model(x, targets=y)
            val_loss += loss.item()
    return val_loss

class MyDataset(Dataset):
    def __init__(self,path,tokenizer, block_size=1024):
        #self.tokenizer = BertTokenizer.from_pretrained("/www/wwwroot/data/jupyter/models/modelscope/hub/tiansz/bert-base-chinese")
        self.tokenizer =tokenizer
        self.block_size = block_size

        # 调整特殊标记
        self.eos_token = self.tokenizer.sep_token_id  # 使用[SEP]作为结束符
        self.pad_token = self.tokenizer.pad_token_id  # 使用[PAD]填充
        
        # 读取数据逻辑保持不变
        self.encoded_data = []
        self.max_lines = 800000

        with open(path, 'r', encoding='utf-8') as f:
            try:
                # 尝试解析完整JSON数组
                data = json.load(f)
                #texts = [item["text"] for item in data[:self.max_lines]]
                texts = [item["text"] for item in data]
            except json.JSONDecodeError:
                # 回退到逐行解析
                f.seek(0)
                texts = []
                for i, line in enumerate(f):
                    if i >= self.max_lines:
                        break
                    try:
                        texts.append(json.loads(line)["text"])
                    except:
                        continue


        # 编码处理（适配BERT）
        full_encoded = []
        j=0
        for text in texts:
            j += 1
            encoded = self.tokenizer.encode(
                text,
                add_special_tokens=False,  # 禁用自动添加[CLS]/[SEP]
                truncation=True,
                max_length=self.block_size-1  # 为EOS保留位置
            )
            full_encoded.extend(encoded + [self.eos_token])

        print("J:",j)
        # 分块处理（使用PAD填充）
        for i in range(0, len(full_encoded), self.block_size):
            chunk = full_encoded[i:i+self.block_size+1]
            if len(chunk) < self.block_size + 1:
                chunk += [self.pad_token] * (self.block_size + 1 - len(chunk))
            self.encoded_data.append(chunk)



    def __len__(self):
        return len(self.encoded_data)
    
    def __getitem__(self, idx):
        chunk = self.encoded_data[idx]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y

    def encode(self, text):
        """将文本编码为token IDs"""
        return self.tokenizer.encode(text, add_special_tokens=False)

    def decode(self, ids):
        """将token IDs解码为文本"""
        return self.tokenizer.decode(ids, skip_special_tokens=True)

def train_test():
    tokenizer = BertTokenizer.from_pretrained("/www/wwwroot/data/jupyter/models/modelscope/hub/tiansz/bert-base-chinese")
    train_data = MyDataset('/www/wwwroot/data/jupyter/aicode/rag/milvus/mytest/xijanfengbao2.json',tokenizer)
    
    # split traindataset to train and val
    train_dataset, val_dataset = torch.utils.data.random_split(train_data, [0.9, 0.1])
    
    train_loader = DataLoader(train_dataset, batch_size=12, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=12, shuffle=False)
    
    model = GPT(GPTConfig())
    device = "cuda" if torch.cuda.is_available() else "cpu"
    #device="cpu"
    model = model.to(device)
    # 打印模型一共有多少参数
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params / 1e6} M")
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)
    # 设置 cosine 学习率
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)
    
    
    for epoch in range(20):
        train_loss = train(model, optimizer, scheduler, train_loader, val_loader, device)
        val_loss = eval(model, val_loader, device)
        print(f'Epoch: {epoch}, Train Loss: {train_loss/len(train_loader):.4f}, Val Loss: {val_loss/len(val_loader):.4f}')
    
        # 保存模型
        if (epoch<30):
        #if (epoch==9 or epoch==0 or epoch==5):
            avg_val_loss = val_loss / len(val_loader)
            # 保存检查点时包含tokenizer
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'val_loss': avg_val_loss,
                'tokenizer': tokenizer,  # 新增tokenizer保存
                'config': GPTConfig()  # 保存模型配置
            }
            # 保存每个epoch的模型
            torch.save(checkpoint, f'tonepart/model_epoch_{epoch}.pt')
            print(f'save:tonepart/model_epoch_{epoch}.pt')
            
def generate_text(prompt, model, tokenizer, device,max_new_tokens=1024, temperature=1.0, top_k=50):
    # 直接使用tokenizer编码
    input_ids = tokenizer.encode(prompt, add_special_tokens=False)
    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)
    
    # 生成过程...
    generated = model.generate(input_tensor, max_new_tokens)
    
    # 使用tokenizer解码
    return tokenizer.decode(generated[0].cpu().numpy().tolist(), skip_special_tokens=True)
    
    
def generated_test():
    # 重新实例化模型结构（必须与训练时一致）
    # 加载检查点（包含tokenizer）
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # 修改前
    #checkpoint = torch.load('tonepart/model_epoch_0.pt', map_location=device)

    # 修改后（添加 weights_only=False）
    checkpoint = torch.load(
        'tonepart/model_epoch_3.pt', 
        map_location=device,
        weights_only=False  # 显式关闭安全模式
    )
    
    # 重建模型
    #model = GPT(GPTConfig()).to(device)
    model = GPT(checkpoint['config']).to(device)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    # 获取保存的tokenizer
    #tokenizer = BertTokenizer.from_pretrained("/www/wwwroot/data/jupyter/models/modelscope/hub/tiansz/bert-base-chinese")
    tokenizer = checkpoint['tokenizer']

    # 打印前 100 个词汇及其对应的 ID
    '''
    for i, (token, token_id) in enumerate(tokenizer.vocab.items()):
        print(f"ID: {token_id} → Token: {token}")
    '''
    # 生成示例
    prompt = "高耸"
    generated_text = generate_text(prompt, model, tokenizer, device)
    print("生成结果：\n", generated_text)
# 主流程
if __name__ == "__main__":
    #train_test()
    generated_test()


